from timm.models.vision_transformer import Attention, LayerScale
from torch import nn
from util.LoRALayer import LoRALayer
from timm.layers import PatchEmbed, Mlp, DropPath, AttentionPoolLatent
from typing import Any, Callable, Dict, Optional, Sequence, Set, Tuple, Type, Union, List
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from torch.jit import Final

import torch
import torch.nn as nn
from typing import Optional


class Block(nn.Module):
    """
    Transformer Block，包含自注意力层（Multi-Head Self-Attention）和前馈神经网络（MLP）。
    本模块包括以下技术：层归一化（LayerNorm）、DropPath（随机丢弃路径）、LayerScale（层尺度）等。
    """

    def __init__(
            self,
            dim: int,  # 输入特征的维度
            num_heads: int,  # 注意力头的数量
            mlp_ratio: float = 4.,  # MLP 隐藏层的维度与输入维度的比率
            qkv_bias: bool = False,  # 是否在 Q、K、V 中添加偏置
            qk_norm: bool = False,  # 是否对 Q 和 K 进行归一化
            proj_drop: float = 0.,  # 投影矩阵的 dropout 概率
            attn_drop: float = 0.,  # 注意力矩阵的 dropout 概率
            init_values: Optional[float] = None,  # LayerScale 的初始化值（可选）
            drop_path: float = 0.,  # DropPath 的概率，控制是否随机丢弃某些路径
            act_layer: nn.Module = nn.GELU,  # 激活函数层，默认为 GELU
            norm_layer: nn.Module = nn.LayerNorm,  # 归一化层，默认为 LayerNorm
            mlp_layer: nn.Module = Mlp,  # MLP 网络模块，默认为 `Mlp`
    ) -> None:
        """
        初始化 Transformer Block，包含自注意力层、MLP 和可选的 LayerScale、DropPath 等辅助层。

        参数：
            dim: 输入特征的维度。
            num_heads: 注意力头的数量。
            mlp_ratio: 隐藏层大小与输入特征大小的比值，默认为 4。
            qkv_bias: 是否在 Q、K、V 上使用偏置，默认为 False。
            qk_norm: 是否对 Q 和 K 进行归一化，默认为 False。
            proj_drop: 投影矩阵的 Dropout 比率，默认为 0。
            attn_drop: 注意力矩阵的 Dropout 比率，默认为 0。
            init_values: LayerScale 的初始化值，默认为 None。
            drop_path: DropPath 的比率，控制是否丢弃部分路径，默认为 0。
            act_layer: 激活函数层，默认为 GELU。
            norm_layer: 归一化层，默认为 LayerNorm。
            mlp_layer: MLP 层模块，默认使用 `Mlp`。
        """
        super().__init__()

        # 第一层归一化
        self.norm1 = norm_layer(dim)

        # 自注意力机制（Multi-Head Self-Attention）
        self.attn = Attention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            norm_layer=norm_layer,
        )

        # 可选的 LayerScale 层，用于稳定训练过程（如果提供了初始化值）
        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()

        # 可选的 DropPath 层，用于正则化，防止过拟合
        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        # 第二层归一化
        self.norm2 = norm_layer(dim)

        # MLP 层，用于进一步处理特征
        self.mlp = mlp_layer(
            in_features=dim,
            hidden_features=int(dim * mlp_ratio),  # 隐藏层特征维度（输入维度的 mlp_ratio 倍）
            act_layer=act_layer,  # 激活函数层
            drop=proj_drop,  # 投影 Dropout 概率
        )

        # 第二个 LayerScale 层，用于稳定训练过程
        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()

        # 第二个 DropPath 层，用于防止过拟合
        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播函数，通过自注意力和 MLP 进行计算，并添加残差连接（skip connection）。

        参数：
            x (Tensor): 输入张量，形状为 (batch_size, dim)，即一个批次的数据。

        返回：
            Tensor: 输出张量，形状与输入相同，经过自注意力和 MLP 操作后的结果。
        """
        # 1. 经过自注意力计算 + 残差连接
        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))

        # 2. 经过 MLP 计算 + 残差连接
        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))

        return x  # 返回经过自注意力和 MLP 处理后的输出

